/*
 * Equilibrator.cpp
 *
 *  Created on: 2 Aug 2011
 *      Author: Allan
 */

#include "Equilibrator.h"

// C++ includes
#include <algorithm>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <set>
#include <tuple>

// GeoReact includes
#include "../CanonicalReactionSystem.h"
#include "../Multiphase.h"
#include "../PrimarySelector.h"
#include "../Types.h"
#include "../Utilities/Assert.h"

Equilibrator::Equilibrator()
{}

void Equilibrator::AddCondition(const string& type)
{
	conditions.push_back(std::make_tuple(type, "", 0.0));
}

void Equilibrator::AddCondition(const string& type, double value)
{
	conditions.push_back(std::make_tuple(type, "", value));
}

void Equilibrator::AddCondition(const string& type, const string& species, double value)
{
	conditions.push_back(std::make_tuple(type, species, value));
}

void Equilibrator::ResetConditions()
{
	conditions.clear();
}

void Equilibrator::Equilibrate(System& system)
{
	SetupPrimarySpecies(system);
	
	SetupInitialGuess(system);
	
	SetupEquilibriumSolver(system);
	
	  double T = system.GetTemperature();
	  double P = system.GetPressure();
	VectorXd n = system.GetComposition();
	
	auto diagnostic = eSolver.Solve(T, P, n);
	
	if(!diagnostic.converged)
		cerr << "EquilibrateError: The equilibrium calculation did not converge." << endl;
	
	system.SetComposition(n);
}

void Equilibrator::SetupPrimarySpecies(System& system)
{
	// An auxiliar set to hold unique primary species
	set<string> set;
	
	// Loop over all the specified equilibrium conditions
	BOOST_FOREACH(auto condition, conditions)
	{
		const string type    = get<0>(condition);
		const string species = get<1>(condition);
		
		if(type == "pH")
			set.insert("H[+]");
		
		else if(type == "FixedMass")
			set.insert(species);
		
		else if(type == "FixedTotalMass")
			set.insert(species);
		
		else if(type == "FixedActivity")
			set.insert(species);
		
		else if(type == "FixedPartialPressure")
			set.insert(species);
	}
	
	// Create a primary selector object to determine the most frequent species in the reaction system
	PrimarySelector selector(system.GetEquilibriumReactions());
	
	// The set of primary species determined as the most frequent species in the reaction system
	const vector<string> frequent = selector.SelectByFrequency();

	// An iterator to iterate over the most frequent species in the reaction system
	vector<string>::const_iterator iter = frequent.begin();
	
	// Loop and insert some of the most frequent species in the set of primary species,
	// until the required number of primary species has been achieved 
	while(set.size() < frequent.size())
		set.insert(*iter++);
	
	// Set the vector of primary species
	vector<string> jSpecies(set.begin(), set.end());
	
	Assert(conditions.size() == jSpecies.size(), 
			"The number of provided equilibrium conditions "
			"does not match the number of primary species.");
	
	system.SetPrimarySpecies(jSpecies);
}

void Equilibrator::SetupInitialGuess(System& system)
{
	// Initialiaze (n) with 1.0E-7 for all species
	VectorXd n = system.GetComposition();
	
	// Loop over all the specified equilibrium conditions
	BOOST_FOREACH(auto condition, conditions)
	{
		const string type    = get<0>(condition);
		const string species = get<1>(condition);
		const double value   = get<2>(condition);
		
		if(type == "pH")
			n[system["H[+]"]] = pow(10.0, -value);
		
		else if(type == "FixedMass" || type == "FixedActivity" || type == "FixedPartialPressure")
			n[system[species]] = value;
		
		else if(type == "FixedTotalMass")
			n[system[species]] = (value > 0.0) ? value : 1.0E-7;
	}
	
	system.SetComposition(n);
}

void Equilibrator::SetupEquilibriumSolver(const System& system)
{
	eSolver.Initialiaze(system.GetMultiphase(), system.GetEquilibriumReactions());
	
	/// The equilibrium conditions as EquilibriumCondition objects
	vector<EquilibriumCondition> eConditions;
	
	// Loop over all the specified equilibrium conditions
	BOOST_FOREACH(auto condition, conditions)
	{
		const string type    = get<0>(condition);
		const string species = get<1>(condition);
		const double value   = get<2>(condition);
		
		if(type == "ChargeBalanced")
			eConditions.push_back(eSolver.ChargeBalanced());
		
		else if(type == "pH")
			eConditions.push_back(eSolver.FixedAcidity(value));
		
		else if(type == "FixedMass")
			eConditions.push_back(eSolver.FixedMass(species, value));
		
		else if(type == "FixedTotalMass")
			eConditions.push_back(eSolver.FixedTotalMass(species, value));
		
		else if(type == "FixedActivity")
			eConditions.push_back(eSolver.FixedActivity(species, value));
		
		else if(type == "FixedPartialPressure")
			eConditions.push_back(eSolver.FixedPartialPressure(species, value));
		
		else
		{ cerr << "EquilibratorError: The equilibrium condition " << type << " is not valid." << endl; exit(0); }
	}
	
	// Set the equilibrium conditions in the equilibrium solver
	eSolver.SetEquilibriumConditions(eConditions);
}

